{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.image as img\n",
    "import scipy.io as sio\n",
    "import h5py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define functions for generating a synthetic noise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_noise(image, shot_noise=0.01, read_noise=0.0005):\n",
    "  \"\"\"Adds random shot (proportional to image) and read (independent) noise.\"\"\"\n",
    "  variance = image * shot_noise + read_noise\n",
    "  sigma=np.sqrt(variance)\n",
    "  noise=sigma *np.random.normal(0,1,(np.shape(image)[0],np.shape(image)[1]))\n",
    "\n",
    "  out      = image + noise\n",
    "\n",
    "  out=np.maximum(0.0,np.minimum(out,1.0))\n",
    "\n",
    "  return out.astype(np.float32)\n",
    "\n",
    "def random_noise_levels():\n",
    "  \"\"\"Generates random noise levels from a log-log linear distribution.\"\"\"\n",
    "  log_min_shot_noise = np.log(0.0001)\n",
    "  log_max_shot_noise = np.log(0.012)\n",
    "  log_shot_noise     = np.random.uniform(log_min_shot_noise, log_max_shot_noise)\n",
    "\n",
    "  shot_noise = np.exp(log_shot_noise)\n",
    "\n",
    "  line = lambda x: 2.18 * x + 1.20\n",
    "  n= np.random.normal(0,0.26)\n",
    "  log_read_noise = line(log_shot_noise) + n\n",
    "  read_noise     = np.exp(log_read_noise)\n",
    "  return shot_noise, read_noise"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Set parameters for the noise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha = 0.01\n",
    "beta = 0.02 ## beta == sigma in the paper"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Load the FiveK dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_name='./Fivek_train_set.hdf5'\n",
    "tr_data= h5py.File(file_name,\"r\")\n",
    "tr_clean_array=tr_data[\"clean_images\"]\n",
    "print(np.shape(tr_clean_array))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_name='./Fivek_test_set.mat'\n",
    "te_data= sio.loadmat(file_name)\n",
    "te_clean_array=te_data[\"clean_images\"]\n",
    "print(np.shape(te_clean_array))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Generate synthtic datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tr_noisy_array=np.zeros(tr_clean_array.shape)\n",
    "for i in range(0, tr_clean_array.shape[0]):\n",
    "    tr_noisy_array[i] = add_noise(tr_clean_array[i], alpha, sigma**2)\n",
    "    \n",
    "te_noisy_array=np.zeros(te_clean_array.shape)\n",
    "for i in range(0, te_clean_array.shape[0]):\n",
    "    te_noisy_array[i] = add_noise(te_clean_array[i], alpha, sigma**2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Save datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f1 = h5py.File(\"./train_fivek_rawRGB_25000x256x256_cropped_alpha_\"+str(alpha)+\"_beta_\"+str(beta)+\".hdf5\", \"w\")\n",
    "f1.create_dataset(\"noisy_images\", tr_noisy_array.shape, dtype='f', data=tr_noisy_array)\n",
    "f1.create_dataset(\"clean_images\", tr_clean_array.shape, dtype='f', data=tr_clean_array)\n",
    "f1.close()\n",
    "\n",
    "f2 = h5py.File(\"./test_fivek_rawRGB_alpha_\"+str(alpha)+\"_beta_\"+str(beta)+\".hdf5\", \"w\")\n",
    "f2.create_dataset(\"noisy_images\", te_noisy_array.shape, dtype='f', data=te_noisy_array)\n",
    "f2.create_dataset(\"clean_images\", te_clean_array.shape, dtype='f', data=te_clean_array)\n",
    "f2.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
